#include "utils.h"
#include "kpabe_user.h"
#include "streamforce_crypto_KPAbe.h"
#include "string.h"
#include <stdlib.h>
#include <stdio.h>
#include <iostream>

using namespace std;

//kpabe_prv_t* transformKey; 
//element_t userDecKey; //, userMask;
//element_t detKeyK2; //X2
//element_t secretDetKey; //this is S, used in the joining protocol

unordered_map<int,dLogMap> user_discreteLogMap;
unordered_map<int,ptMap> user_plaintextMap;
unordered_map<int,kpabe_pub_t*> userPKMap;
unordered_map<int,kpabe_prv_t*> userTKMap;
unordered_map<int,element_wrap> userDetKeyMap;
unordered_map<int,element_wrap> userDecKeyMap;
unordered_map<int,int> wsUserMap;

unordered_map<int,element_wrap> user_windowKeyMap;
unordered_map<int, ptMap> user_currentRSum;

//unordered_map<int, kpabe_pub_t*> userPubKeyMap;
//unordered_map<int, kpabe_prv_t*> userTKMap;
//unordered_map<int, element_wrap> userDecKeyMap;
//unordered_map<int, element_wrap> userDetKeyMap;

bool joiningEnabled=false; 

int initUser_simple(int userId, int columnId, unsigned char* pubBytes, int pubLen, unsigned char* prvBytes, int privLen, unsigned char*
dec, int decLen, unsigned char *det, int detKeyLen){
	pbc_random_set_deterministic(userId);

	kpabe_pub_t* pub = convertToPubKey(pubBytes, pubLen);
	userPKMap[userId] = pub;
	user_initDiscreteLogMap(50000,userId,pub);
	userTKMap[userId] = convertToPrvKey(pub, prvBytes, privLen);
	element_wrap udk; 
	element_init_Zr(udk.element, pub->p);
	element_from_bytes(udk.element, dec);
	userDecKeyMap[userId] = udk;
	if (detKeyLen!=0){
		joiningEnabled = true;
		element_wrap tmp;
		element_init_Zr(tmp.element, pub->p);
		element_from_bytes(tmp.element, det);
		userDetKeyMap[userId] = tmp;	
	}

	for (int i=0; i<columnId; i++){
		element_init_GT(user_currentRSum[userId].map[i].element, pub->p);
		element_set_si(user_currentRSum[userId].map[i].element,1);	
	}	
	return 0;
}

JNIEXPORT void JNICALL Java_streamforce_crypto_KPAbe_initUser_1simple
  (JNIEnv *env, jobject obj, jint userId, jint columnId, jbyteArray pubB, jbyteArray tkB, jbyteArray decB, jbyteArray detB){
	int pl = env->GetArrayLength(pubB);
	int tkl = env->GetArrayLength(tkB);
	int dl = env->GetArrayLength(decB);
	int detl = env->GetArrayLength(detB);

	unsigned char *pubBytes = (unsigned char*)malloc(pl*sizeof(char));
	unsigned char *prvBytes = (unsigned char*)malloc(tkl*sizeof(char));
	unsigned char *decBytes = (unsigned char*)malloc(dl*sizeof(char));
	unsigned char *detBytes = (unsigned char*)malloc(detl*sizeof(char));
	
	env->GetByteArrayRegion(pubB,0,pl,(jbyte*)pubBytes);
	env->GetByteArrayRegion(tkB,0,tkl,(jbyte*)prvBytes);
	env->GetByteArrayRegion(decB,0,dl,(jbyte*)decBytes);
	env->GetByteArrayRegion(detB,0,detl,(jbyte*)detBytes);	

	initUser_simple(userId,columnId,pubBytes,pl,prvBytes,tkl,decBytes,dl,detBytes,detl);
}


int initUser_slidingWindow_Aggregate1(int userId, unsigned char* pubBytes, int pubLen, unsigned char* prvBytes, int privLen, unsigned char*
dec, int decLen, int windowSize, unsigned char *windowKey){
	kpabe_pub_t* pub = convertToPubKey(pubBytes, pubLen);
	userPKMap[userId] = pub;
	user_initDiscreteLogMap(50000,userId,pub);
	userTKMap[userId] = convertToPrvKey(pub,prvBytes, privLen);
	wsUserMap[userId] = windowSize; 

	element_wrap tmp; 
	element_init_Zr(tmp.element, pub->p);
	element_from_bytes(tmp.element, dec);
	userDecKeyMap[userId]= tmp; 

	element_wrap tmp1;
	element_init_Zr(tmp1.element,pub->p);
	element_from_bytes(tmp1.element,windowKey);
	user_windowKeyMap[userId] = tmp1;

	return 0;
}

JNIEXPORT void JNICALL Java_streamforce_crypto_KPAbe_initUser_1slidingWindow_1Aggregate1
  (JNIEnv *env, jobject obj, jint userId, jbyteArray pubB, jbyteArray prvB, jbyteArray decB,
jint windowSize, jbyteArray windowKey){
	int pl = env->GetArrayLength(pubB);
	int tkl = env->GetArrayLength(prvB);
	int dl = env->GetArrayLength(decB);
	int wl = env->GetArrayLength(windowKey);
	
	unsigned char *pubBytes = (unsigned char*)malloc(pl*sizeof(char));
	unsigned char *prvBytes = (unsigned char*)malloc(tkl*sizeof(char));
	unsigned char *decBytes = (unsigned char*)malloc(dl*sizeof(char));
	unsigned char *windowBytes = (unsigned char*)malloc(wl*sizeof(char));	

	env->GetByteArrayRegion(pubB,0,pl,(jbyte*)pubBytes);
	env->GetByteArrayRegion(prvB,0,tkl,(jbyte*)prvBytes);
	env->GetByteArrayRegion(decB,0,dl,(jbyte*)decBytes);
	env->GetByteArrayRegion(windowKey,0,wl,(jbyte*)windowBytes);

	initUser_slidingWindow_Aggregate1(userId, pubBytes,pl,prvBytes,tkl,decBytes,dl,windowSize,
windowBytes);

}

//received the ciphertext, need to umask the e(g,g)^{y.s} first
//then perform discrete log.
//the input cipher has been transformed
int decrypt_simple(int userId, unsigned char* cipher, unsigned char* originalCipher){
	element_t userCipher, val;
	kpabe_pub_t* pub = userPKMap[userId];
	
	element_init_GT(userCipher, pub->p);
	element_init_GT(val, pub->p);
	element_from_bytes(userCipher, cipher);
	element_from_bytes(val, originalCipher);	

	element_t decVal;
	kpabe_dec_proxy(pub,userCipher,(userDecKeyMap[userId]).element,val,decVal);
	return toPlaintext_simple(userId, decVal); 		
}

JNIEXPORT jint JNICALL Java_streamforce_crypto_KPAbe_decrypt_1simple
  (JNIEnv *env, jobject obj, jint userId, jbyteArray cipher, jbyteArray originalCipher){
	int cl = env->GetArrayLength(cipher);
	int ocl = env->GetArrayLength(originalCipher);
	unsigned char *cipherBytes = (unsigned char*)malloc(cl*sizeof(char));
	unsigned char *originalCipherBytes = (unsigned char*)malloc(ocl*sizeof(char));
	
	env->GetByteArrayRegion(cipher,0,cl,(jbyte*)cipherBytes);
	env->GetByteArrayRegion(originalCipher,0,ocl,(jbyte*)originalCipherBytes);
	int val = decrypt_simple(userId, cipherBytes, originalCipherBytes);
	return val;
}

int decrypt_slidingWindow_Aggregate3(int userId, int rUserId, int columnId, unsigned char* userCipher, unsigned char* originalCipher, unsigned char *rUserCipher, unsigned
char* rOriginalCipher){
	element_t uc, oc, ruc, roc;
	kpabe_pub_t* pub = userPKMap[userId];
	element_init_GT(uc, pub->p);
	element_from_bytes(uc, userCipher);
	element_init_GT(oc, pub->p);
	element_from_bytes(oc, originalCipher);
	element_init_GT(ruc, pub->p);
	element_from_bytes(ruc, rUserCipher);
	element_init_GT(roc, pub->p);
	element_from_bytes(roc, rOriginalCipher);
	
	element_t decVal1;
	element_init_GT(decVal1, pub->p);
	element_t tmp;
	element_init_GT(tmp, pub->p);
	element_pow_zn(tmp,uc,userDecKeyMap[userId].element);
	element_div(decVal1,oc,tmp);

	element_t decVal2;
	element_init_GT(decVal2, pub->p);
	element_pow_zn(tmp,ruc,userDecKeyMap[rUserId].element);
	element_div(decVal2,roc,tmp);
	
	element_t temp_result;
	element_init_GT(temp_result, pub->p);
	element_div(temp_result,decVal1,decVal2);
	
	element_t result;element_init_GT(result,pub->p);
	element_mul(result,temp_result,user_currentRSum[rUserId].map[columnId].element);
	element_set(user_currentRSum[rUserId].map[columnId].element,decVal2);
	
	return toPlaintext_simple(userId,result);	
}

//compute the unmask value, before doing the same as decrypt_simple
int decrypt_slidingWindow_Aggregate1(int userId, unsigned char* cipher, unsigned char* originalCipher){
	element_t userCipher, val;
	kpabe_pub_t* pub = userPKMap[userId];
	
	element_wrap wk = user_windowKeyMap[userId]; 
	element_init_GT(userCipher, pub->p);
	element_init_GT(val, pub->p);
	element_from_bytes(userCipher, cipher);
	element_from_bytes(val, originalCipher);	
	
	element_t decVal;
	kpabe_dec_proxy(pub,userCipher,(userDecKeyMap[userId]).element,val, decVal);
	return toPlaintext_slidingWindow(userId, decVal,wk.element); 		
}

JNIEXPORT jint JNICALL Java_streamforce_crypto_KPAbe_decrypt_1slidingWindow_1Aggregate3
  (JNIEnv *env, jobject obj, jint userId, jint rUserId, jint columnId, jbyteArray userCipher, jbyteArray originalCipher, jbyteArray rUserCipher, jbyteArray rOriginalCipher){
	int cl = env->GetArrayLength(userCipher);
	int ocl = env->GetArrayLength(originalCipher);
	unsigned char *ucBytes = (unsigned char*)malloc(cl*sizeof(char));
	unsigned char *ocBytes = (unsigned char*)malloc(ocl*sizeof(char));
	unsigned char *rucBytes = (unsigned char*)malloc(cl*sizeof(char));
	unsigned char *rocBytes = (unsigned char*)malloc(ocl*sizeof(char));
	env->GetByteArrayRegion(userCipher,0,cl,(jbyte*)ucBytes);
	env->GetByteArrayRegion(originalCipher,0,ocl,(jbyte*)ocBytes);
	env->GetByteArrayRegion(rUserCipher,0,cl,(jbyte*)rucBytes);
	env->GetByteArrayRegion(rOriginalCipher,0,ocl,(jbyte*)rocBytes);
	return decrypt_slidingWindow_Aggregate3(userId,rUserId, columnId, ucBytes,ocBytes,rucBytes,rocBytes);
}


JNIEXPORT jint JNICALL Java_streamforce_crypto_KPAbe_decrypt_1slidingWindow_1Aggregate1
  (JNIEnv *env, jobject obj, jint userId, jbyteArray cipher, jbyteArray originalCipher){
	int cl = env->GetArrayLength(cipher);
	int ocl = env->GetArrayLength(originalCipher);
	unsigned char *cipherBytes = (unsigned char*)malloc(cl*sizeof(char));
	unsigned char *originalCipherBytes = (unsigned char*)malloc(ocl*sizeof(char));
	
	env->GetByteArrayRegion(cipher,0,cl,(jbyte*)cipherBytes);
	env->GetByteArrayRegion(originalCipher,0,ocl,(jbyte*)originalCipherBytes);
	return decrypt_slidingWindow_Aggregate1(userId, cipherBytes, originalCipherBytes);

}



JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_getJoiningKey
  (JNIEnv *env, jobject obj, jint userId1, jint userId2){
	unsigned char *joinKey1, *joinKey2;
	int lenKey1, lenKey2;
	computeJoinKey(userId1, userId2, &joinKey1, &lenKey1, &joinKey2, &lenKey2);
	jbyteArray ret = env->NewByteArray(lenKey1+lenKey2);
	env->SetByteArrayRegion(ret,0,lenKey1,(jbyte*)joinKey1);
	env->SetByteArrayRegion(ret,lenKey1,lenKey2,(jbyte*)joinKey2);
	return ret;	 
}

//for joining 2 streams, compute s/x1, s/x2
int computeJoinKey(int userId1, int userId2, unsigned char **cipher1, int *len1, unsigned
char **cipher2, int *len2){
	kpabe_pub_t* pub1 = userPKMap[userId1];
	kpabe_pub_t* pub2 = userPKMap[userId2];
	element_wrap x1 = userDetKeyMap[userId1];
	element_wrap x2 = userDetKeyMap[userId2];
	//generate random s
	element_t s;
	element_init_Zr(s,pub1->p);
	element_random(s);

	//compute s/x1
	element_t val;
	element_init_Zr(val,pub1->p);
	element_div(val,s,x1.element);
	*len1 = element_length_in_bytes(val);
	*cipher1 = (unsigned char*)malloc((*len1)*sizeof(char));
	element_to_bytes(*cipher1,val);

	element_init_Zr(val,pub1->p);	
	element_div(val,s,x2.element);
	*len2 = element_length_in_bytes(val);
	*cipher2 = (unsigned char*)malloc((*len2)*sizeof(char));
	element_to_bytes(*cipher2,val);
	return 0;	
} 
